from model_utils.yolov1_backbone_improve import YOLOv1Net
from proposal_utils.yolov1_bbox import dst_bbox
import torch
from torch.utils.data import DataLoader, Dataset
from torch import nn
import cv2
import numpy as np

device = torch.device("cuda:0")
model = YOLOv1Net(2)
model = model.to(device)
batch_size = 20


class CustomDataset(Dataset):
    def __init__(self, data):
        super().__init__()
        self.data = data
        self.features = data['features']
        self.ground_truths = data['ground_truths']
        self.cls_list = data['cls_list']

    def __getitem__(self, index):
        return self.features[index], self.ground_truths[index], self.cls_list[index]

    def __len__(self):
        return len(self.features)


train_datasets = CustomDataset(torch.load("./VOCdevkit/cat_dog_train_data.pth", weights_only=False))
train_dataloader = DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
valid_datasets = CustomDataset(torch.load("./VOCdevkit/cat_dog_valid_data.pth", weights_only=False))
valid_dataloader = DataLoader(valid_datasets, batch_size=batch_size)


def loss_bbox(box1, box2):
    # 中心点位置损失
    box1_xy = box1[:, :, :, [0, 1, 5, 6]]
    box2_xy = box2[:, :, :, [0, 1, 5, 6]]
    loss_xy = torch.mean((box1_xy - box2_xy) ** 2)
    # 尺度损失
    box1_wh = box1[:, :, :, [2, 3, 7, 8]]
    box2_wh = box2[:, :, :, [2, 3, 7, 8]]
    loss_wh = torch.mean((box1_wh ** 0.5 - box2_wh ** 0.5) ** 2)
    loss_loc = loss_xy + loss_wh
    # 置信度损失
    loss_conf1 = torch.mean((box1[:, :, :, [4, 9]] - box2[:, :, :, [4, 9]]) ** 2)
    loss_conf2 = torch.mean((1 - (box1[:, :, :, [4, 9]]) - (1 - box2[:, :, :, [4, 9]])) ** 2)
    print(loss_conf1,loss_conf2)
    # 类别损失
    loss_cls = torch.mean((box1[:, :, :, -2:] - box2[:, :, :, -2:]) ** 2)
    return 5 * loss_loc + loss_conf1 + 0.5 * loss_conf2 + loss_cls


sgd = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
model.train()
epochs = 100
for epoch in range(epochs):
    losses = []
    for image, ground_truth, cls in train_dataloader:
        sgd.zero_grad()
        predict_bbox = model((image / 255.).to(device).float())
        real_bbox = dst_bbox(ground_truth, cls, 2, 2)
        loss = loss_bbox(predict_bbox, real_bbox.to(device))
        loss.backward()
        losses.append(loss.item())
        sgd.step()

    epoch_loss = np.mean(losses)
    print("epoch:", epoch + 1, "loss:", epoch_loss)
